import sys
sys.path.append("..")
import pyrtl
from  pyrtl import GPUSim
import or1200_definitions
from functools import reduce

class Or1200Pic(object):
    def __init__(self):
        self.rst = pyrtl.Input(bitwidth=1, name='rst')
        self.spr_cs = pyrtl.Input(bitwidth=1, name='spr_cs')
        self.spr_write = pyrtl.Input(bitwidth=1, name='spr_write')
        self.spr_addr = pyrtl.Input(bitwidth=32, name='spr_addr')
        self.spr_dat_i = pyrtl.Input(bitwidth=32, name='spr_dat_i')

        self.pic_wakeup = pyrtl.Output(bitwidth=1, name='pic_wakeup')
        self.intr = pyrtl.WireVector(bitwidth=1, name='intr')

        self.pic_int = pyrtl.Input(bitwidth=or1200_definitions.OR1200_PIC_INTS, name='pic_int')
        self.picmr = pyrtl.Register(bitwidth=or1200_definitions.OR1200_PIC_INTS-2, name='picmr')
        self.picsr = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_PIC_INTS, name='picsr')

        self.picmr_sel = pyrtl.WireVector(bitwidth=1, name='picmr_sel')
        self.picsr_sel = pyrtl.WireVector(bitwidth=1, name='picsr_sel')
        self.um_ints = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_PIC_INTS, name='um_ints')
        self.spr_dat_o = pyrtl.Register(bitwidth=32, name='spr_dat_o')

        self.picmr_sel <<= pyrtl.select(self.spr_cs.__and__(self.spr_addr[or1200_definitions.OR1200_PICOFS_BITS]\
                                                          .__eq__(or1200_definitions.OR1200_PIC_OFS_PICMR)),\
                                      pyrtl.Const(0b1, bitwidth=1), pyrtl.Const(0b0, bitwidth=1))
        self.picsr_sel <<= pyrtl.select(self.spr_cs.__and__(self.spr_addr[or1200_definitions.OR1200_PICOFS_BITS] \
                                                          .__eq__(or1200_definitions.OR1200_PIC_OFS_PICMR)), \
                                      pyrtl.Const(0b1, bitwidth=1), pyrtl.Const(0b0, bitwidth=1))

        # initialize all the parts
        write2Picmr = Write2Picmr()
        write2Picsr = Write2Picsr()
        readPic = ReadPic()

        # establish connection relations for Write to PICMR
        self.picmr.next <<= write2Picmr.picmr
        write2Picmr.rst <<= self.rst
        write2Picmr.picmr_sel <<= self.picmr_sel
        write2Picmr.spr_write <<= self.spr_write
        write2Picmr.spr_dat_i <<= self.spr_dat_i

        # establish connection relations for Write to PICSR, both CPU and external ints
        self.picsr <<= write2Picsr.picsr
        write2Picsr.rst <<= self.rst
        write2Picsr.picsr_sel <<= self.picsr_sel
        write2Picsr.spr_write <<= self.spr_write
        write2Picsr.spr_dat_i <<= self.spr_dat_i
        write2Picsr.um_ints <<= self.um_ints

        # establish connection relations for Read PIC registers
        self.spr_dat_o.next <<= readPic.spr_dat_o
        #self.um_ints <<= readPic.um_ints
        #self.intr <<= readPic.intr
        #self.pic_wakeup <<= readPic.pic_wakeup
        readPic.spr_addr <<= self.spr_addr
        readPic.picmr <<= self.picmr
        readPic.picsr <<= self.picsr
        #readPic.pic_int <<= self.pic_int

        self.um_ints <<= self.pic_int & pyrtl.concat(self.picmr, pyrtl.Const(0b11, bitwidth=2))
        self.intr <<= or1200_definitions.Or(self.um_ints)
        self.pic_wakeup <<= self.intr

# Write tp PICMR
class Write2Picmr(object):
    def __init__(self):
        self.rst = pyrtl.WireVector(bitwidth=1)
        self.picmr = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_PIC_INTS - 2)
        self.picmr_sel = pyrtl.WireVector(bitwidth=1)
        self.spr_write = pyrtl.WireVector(bitwidth=1)
        self.spr_dat_i = pyrtl.WireVector(bitwidth=32)

        with pyrtl.conditional_assignment:
            with self.rst:
                self.picmr |= pyrtl.concat(pyrtl.Const(0b1, bitwidth=1), \
                                           pyrtl.concat(or1200_definitions.OR1200_PIC_INTS-\
                                                        pyrtl.Const(0b0, bitwidth=3)))
            with self.picmr_sel & self.spr_write:
                self.picmr |= self.spr_dat_i[2:or1200_definitions.OR1200_PIC_INTS]

# Write to PICSR, both CPU and external ints
class Write2Picsr(object):
    def __init__(self):
        self.rst = pyrtl.WireVector(bitwidth=1)
        self.picsr = pyrtl.Register(bitwidth=or1200_definitions.OR1200_PIC_INTS)
        self.picsr_sel = pyrtl.WireVector(bitwidth=1)
        self.spr_write = pyrtl.WireVector(bitwidth=1)
        self.spr_dat_i = pyrtl.WireVector(bitwidth=32)
        self.um_ints = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_PIC_INTS)

        with pyrtl.conditional_assignment:
            with self.rst:
                self.picsr.next |= pyrtl.concat(pyrtl.Const(0b0, bitwidth=or1200_definitions.OR1200_PIC_INTS))
            with self.picsr_sel & self.spr_write:
                self.picsr.next |= self.spr_dat_i[:or1200_definitions.OR1200_PIC_INTS] | self.um_ints
            with pyrtl.otherwise:
                self.picsr.next |= self.picsr | self.um_ints

# Read PIC registers
class ReadPic(object):
    def __init__(self):
        self.spr_addr = pyrtl.WireVector(bitwidth=32)
        self.picmr = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_PIC_INTS - 2)
        self.picsr = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_PIC_INTS)
        self.spr_dat_o = pyrtl.WireVector(bitwidth=32)
        #self.um_ints = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_PIC_INTS)
        #self.pic_int = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_PIC_INTS)
        #self.intr = pyrtl.WireVector(bitwidth=1)
        #self.pic_wakeup = pyrtl.WireVector(bitwidth=1)

        with pyrtl.conditional_assignment:
            with self.spr_addr[or1200_definitions.OR1200_PICOFS_BITS].__eq__(or1200_definitions.OR1200_PIC_OFS_PICMR):
                self.spr_dat_o |= pyrtl.concat(self.picmr, pyrtl.Const(0b0, bitwidth=2))
                # with self.spr_addr[or1200_definitions.OR1200_PICOFS_BITS].__eq__(or1200_definitions.OR1200_PIC_UNUSED_ZERO):
                #     self.spr_dat_o |= pyrtl.Const(0b0, bitwidth=32)
            with pyrtl.otherwise:
                self.spr_dat_o |= self.picsr
            # with self.spr_addr[or1200_definitions.OR1200_PICOFS_BITS].__eq__(or1200_definitions.OR1200_PIC_UNUSED_ZERO):
            #     self.spr_dat_o |= pyrtl.Const(0b0, bitwidth=32)

            #self.um_ints <<= self.pic_int & pyrtl.concat(self.picmr, pyrtl.Const(0b11, bitwidth=2))
            #self.intr <<= reduce(lambda x, y: x | y, [self.um_ints[i] for i in range(self.um_ints.bitwidth)])
            #self.pic_wakeup <<= self.intr



if __name__ == '__main__':
    or1200Pic = Or1200Pic()
    sim = pyrtl.GPUSim_now.GPUSim(65536)
    sim.create_dll('pic.cu')
    # print(pyrtl.working_block())
    # ff = open("D:\work_doc\PyRTL\examples\Opencore\or1200_pic.c", 'w')
    # essent = pyrtl.ESSENT()
    # essent._create_code(lambda s: ff.write(s + '\n'))
    # ff.close()
    # com = pyrtl.CompiledSimulation()
    #
    # input_width = []
    # from pyrtl.wire import Input
    #
    # IN = pyrtl.working_block().wirevector_subset(Input)
    # for inn in IN:
    #     name = inn.name
    #     width = inn.bitwidth
    #     input_width.append([name, width])
    #
    # import random
    #
    # for i in range(10):
    #
    #     inputs = {}
    #     inputs_old = {}
    #
    #     for inp in input_width:
    #         inputs[inp[0]] = []
    #
    #         result = random.randint(0, 2 ** inp[1] - 1)
    #         result_old = result
    #         inputs_old[inp[0]] = [result, result_old]
    #
    #     for j in range(10000):
    #         for inp in input_width:
    #             if random.randint(0, 20) in [0, 5]:
    #                 result = random.randint(0, 2 ** inp[1] - 1)
    #                 inputs[inp[0]].append(result)
    #                 inputs_old[inp[0]][1] = result
    #             else:
    #                 inputs[inp[0]].append(inputs_old[inp[0]][1])
    #
    #     time1 = essent.step_multiple(inputs)
    #     time2 = com.step_multiple(inputs)
    #
    #     print('==========')
    #     print('essent:', time1)
    #     print('com:', time2)
    #     print(time2 / time1)
